"""Add a private key by mnemonic
"""
from test import mnemonic2
from chia.wallet.puzzles.p2_delegated_puzzle_or_hidden_puzzle import puzzle_for_pk
from chia.util.byte_types import hexstr_to_bytes
from chia.util.bech32m import encode_puzzle_hash, decode_puzzle_hash
from chia.consensus.coinbase import create_puzzlehash_for_pk
from chia.util.keychain import Keychain
from blspy import PrivateKey, G1Element
from chia.wallet.derive_keys import master_sk_to_wallet_sk
from chia.types.blockchain_format.program import Program
from chia.types.blockchain_format.sized_bytes import bytes32

keychain: Keychain = Keychain()


class KeyStore:
    def __init__(self, mnemonic) -> None:
        self.mnemonic = mnemonic
        self.master_sk = self.add_private_key(self.mnemonic)
        self.master_pk = self.master_sk.get_g1()
        self.fingerprint = self.master_pk.get_fingerprint()
        self.wallet_sk = self.create_wallet_sk(index=1)
        self.wallet_pk = self.get_wallet_pk()
        self.puzzle_hash = self.get_puzzlehash_for_pk()
        self.xch_address = self.get_xch_address_for_puzzlehash()

    def add_private_key(self, mnemonic: str, passphrase: str = ""):
        """
        通过助记词获取私钥
        """
        return keychain.add_private_key(mnemonic, passphrase)

    def get_public_key(self, sk: PrivateKey) -> G1Element:
        """通过私钥获取公钥"""
        return sk.get_g1()

    def get_fingerprint(pk: G1Element) -> int:
        """通过公钥获取指纹
        """
        return pk.get_fingerprint()

    def create_wallet_sk(self, index: int = 0) -> PrivateKey:
        """主私钥派生出钱包私钥"""
        sk: PrivateKey = self.master_sk
        return master_sk_to_wallet_sk(sk, index)  # 钱包私钥

    def get_wallet_pk(self) -> G1Element:
        """通过钱包私钥获取钱包公钥"""
        sk = self.wallet_sk
        return sk.get_g1()

    def get_puzzlehash_for_pk(self) -> bytes32:
        """通过公钥获取puzzle_hash"""
        pk = self.wallet_pk
        return create_puzzlehash_for_pk(pk)

    @property
    def puzzlehash_str(self) -> str:
        """获取puzzlehash的hex字符串"""
        return f"0x{bytes(self.puzzle_hash).hex()}"

    # def get_puzzle_reveal_for_pk(pk: G1Element):
    #     """通过公钥获取puzzle_reveal"""
    #     return puzzle_for_pk(pk)

    def get_xch_address_for_puzzlehash(self):
        """通过puzzlehash获取钱包地址"""
        puzzle_hash: str = self.puzzlehash_str
        return encode_puzzle_hash(
            hexstr_to_bytes(puzzle_hash), "xch")

    @staticmethod
    def puzzlehash_to_xchaddress(puzzl_hash: str):
        if not puzzl_hash.startswith('0x') or not puzzl_hash.startswith("0X"):
            puzzl_hash = '0x' + puzzl_hash
        return encode_puzzle_hash(hexstr_to_bytes(puzzl_hash), "xch")

    def get_puzzlehash_for_xchaddress(xch_address: str):
        """通过钱包地址获取puzzlehash"""
        return decode_puzzle_hash(xch_address).hex()  # puzzle_hash 不带0x


if __name__ == "__main__":
    keystore = KeyStore(mnemonic=mnemonic2)
    wallet_sk = keystore.create_wallet_sk()
    print(
        keystore.wallet_sk,
        keystore.puzzlehash_str,
        keystore.puzzlehash_to_xchaddress(
            'ad4a3863dd2cd11956fd8674a5c7dfc5065155807dd3295bfff9bf84796987dd')
    )
